# -- coding: utf-8 --
# @Time : 2022/6/16
# @Author : ykk648
# @Project : https://github.com/ykk648/AI_power
"""
insert body coco_whole_body 133-kp to openpose 25+21+21-kp
"""
import numpy as np

from mocap_lib.skeleton_transfer.keypoints_map import POSE_BODY_25_BODY_PARTS, COCO_JOINTS, COCO_WHOLE_BODY_JOINTS
from mocap_lib.skeleton_transfer.utils import geJointIndexFromName, computeNeck, computeMidHip


def cocowb_2_openpose(coco_wb_kps):

    # whole body
    left_hand_kps = np.array(coco_wb_kps[91:112])
    right_hand_kps = np.array(coco_wb_kps[112:133])
    coco_kps = coco_wb_kps[:23]

    openpose_25_kps = []
    for i, openposeKeypointName in enumerate(POSE_BODY_25_BODY_PARTS):
        try:
            cocoIndex = geJointIndexFromName(openposeKeypointName, COCO_WHOLE_BODY_JOINTS)
            cocoKeypoint = coco_kps[cocoIndex]

            new_keypoint = [cocoKeypoint[0], cocoKeypoint[1], cocoKeypoint[2]]

            openpose_25_kps.append(new_keypoint)
        except ValueError:
            if openposeKeypointName == "Neck":
                new_keypoint = computeNeck(coco_kps)
                openpose_25_kps.append(new_keypoint)
            elif openposeKeypointName == "MidHip":
                new_keypoint = computeMidHip(coco_kps)
                openpose_25_kps.append(new_keypoint)
            elif openposeKeypointName == "Background":
                pass
            else:
                raise 'should not be here'

    return left_hand_kps, right_hand_kps, np.array(openpose_25_kps)


if __name__ == '__main__':
    kps_in = [[1774.271338224411, 775.3212462961674, 0.9747584462165833], [1808.9348201155663, 747.1897142231464, 0.9500738382339478], [1761.7484991550446, 750.5351687073708, 0.9522068500518799], [1878.387617945671, 764.9470171332359, 1.0229778289794922], [1759.26577013731, 772.4447884857655, 0.8659712672233582], [1953.5112162828445, 949.7534110546112, 0.9093562960624695], [1729.8776029348373, 940.9883849024773, 0.9314572811126709], [1988.6106472611427, 1188.1878989934921, 0.9280189275741577], [1707.3157907128334, 1142.5144288539886, 0.8408532738685608], [1953.2337095737457, 1395.901929140091, 1.0020561218261719], [1706.6597483158112, 1317.3678405284882, 0.852715253829956], [1917.8310883045197, 1363.7788480520248, 0.7779121398925781], [1766.171604692936, 1361.103828072548, 0.7959856390953064], [1942.0537549257278, 1642.6289583444595, 0.9424837231636047], [1804.7740798592567, 1634.1614379882812, 0.8356976509094238], [1945.142244398594, 1904.871351480484, 0.825141966342926], [1852.4495774507523, 1866.5656069517136, 0.8548523187637329], [1910.8716644644737, 2016.1494747400284, 0.6750859618186951], [1913.343696296215, 2013.0685818195343, 0.6916803121566772], [1947.6811982393265, 1919.6125905513763, 0.8010267019271851], [1761.364278972149, 1941.6512742042542, 0.619780957698822], [1763.7608103752136, 1938.7530086040497, 0.6242167353630066], [1860.1595613360405, 1879.3040434122086, 0.6505212783813477], [1757.131293952465, 749.4858367741108, 0.9894435405731201], [1756.6606177687645, 765.0954599678516, 0.9894559383392334], [1757.4108677506447, 780.376183450222, 0.9744463562965393], [1759.446950495243, 795.6795932352543, 0.9435798525810242], [1762.4079005718231, 811.6742376089096, 0.9784952402114868], [1767.1554357409477, 826.4706937968731, 0.9555299282073975], [1774.8719309568405, 838.7099540829659, 0.962067186832428], [1788.7493335008621, 846.4731913506985, 0.9193580150604248], [1805.7499009370804, 844.5342876315117, 0.9364320039749146], [1822.213025212288, 837.8297617137432, 0.9542384743690491], [1836.6845091581345, 829.2915210723877, 0.9378867149353027], [1849.6050561666489, 820.969549626112, 0.988083004951477], [1859.4024382829666, 812.5037604570389, 1.0004665851593018], [1864.198136627674, 797.0032020509243, 0.9767622351646423], [1867.3362879157066, 781.170591622591, 0.9437951445579529], [1869.3835329413414, 764.5907543301582, 0.9613441228866577], [1870.2789442539215, 747.9425222277641, 0.9832771420478821], [1752.5094377994537, 740.491570442915, 0.9415059089660645], [1755.9423043131828, 737.7704026997089, 0.9457069039344788], [1760.576769709587, 737.1886721551418, 0.9508200883865356], [1765.5024362802505, 737.2556458413601, 0.9740623235702515], [1770.7091236114502, 737.9240132570267, 0.972709059715271], [1789.6794719099998, 734.6444213986397, 0.9783438444137573], [1799.8617976903915, 731.8639281988144, 0.9982303380966187], [1810.0418496727943, 730.9591374695301, 0.9765279293060303], [1819.8040912747383, 732.1981248259544, 0.9939391613006592], [1829.601525068283, 736.009527027607, 0.9749893546104431], [1779.1310129761696, 750.1261951327324, 0.985303521156311], [1775.8612397909164, 758.7913550436497, 1.0057941675186157], [1772.8083042502403, 767.4351205825806, 0.9786044359207153], [1770.2868173122406, 776.865826934576, 0.9777886867523193], [1766.7060505747795, 785.8237383365631, 0.9968843460083008], [1771.1867762207985, 787.7917621731758, 0.98822420835495], [1776.8337018489838, 788.1493944227695, 0.9776839017868042], [1782.7296600341797, 786.9892665147781, 0.9600222706794739], [1788.5625203251839, 785.4115090966225, 0.9571734070777893], [1758.0200389027596, 751.7940776348114, 0.9784285426139832], [1761.6748074293137, 749.945350676775, 0.9727488160133362], [1766.5626979470253, 749.7107619047165, 0.9789206981658936], [1771.1939076781273, 751.7554489076138, 0.9724154472351074], [1766.5634731054306, 752.9972526729107, 0.9818798899650574], [1762.173182606697, 753.0550278127193, 0.9850895404815674], [1796.236485183239, 749.6632705330849, 0.9822853803634644], [1803.9784055948257, 746.0865604579449, 0.9733872413635254], [1813.1352484822273, 746.3306061625481, 0.9517993330955505], [1821.5957407355309, 748.9158111214638, 0.9691038131713867], [1813.2113690376282, 751.2487536966801, 0.9608550071716309], [1804.568197786808, 751.2391675710678, 0.982032835483551], [1766.9244385361671, 808.6954330503941, 1.0051645040512085], [1769.5562563538551, 804.292300760746, 1.0011354684829712], [1774.2323735952377, 801.5422712862492, 0.9812349677085876], [1777.3411722183228, 801.1994704008102, 0.9650382399559021], [1781.2991310358047, 800.3028188347816, 0.9422434568405151], [1792.699592411518, 801.5905119776726, 0.9543154835700989], [1803.7415688633919, 805.1192397475243, 0.9403200745582581], [1796.8195076584816, 810.3827978670597, 0.95192950963974], [1788.6099600195885, 814.7508671581745, 0.9304230213165283], [1779.346351981163, 817.4082135260105, 0.9493058919906616], [1774.2015739679337, 816.5669341087341, 0.9740787148475647], [1770.1561772823334, 813.5166857838631, 0.9890342950820923], [1768.0753420591354, 808.8013713657856, 1.003024935722351], [1773.0034374594688, 806.9128529429436, 0.9954053163528442], [1778.6033884882927, 805.7820776998997, 0.969303548336029], [1790.6679022312164, 804.6682525873184, 0.9580917358398438], [1801.7620726823807, 805.2996191084385, 0.9488523006439209], [1791.847383260727, 808.134270042181, 0.9512938261032104], [1779.953094303608, 810.3065481185913, 0.9657453298568726], [1773.7954426407814, 810.3798781037331, 0.9914306998252869], [1950.5041701197624, 1402.8576838970184, 0.9920955300331116], [1929.218682050705, 1425.5850182771683, 0.9500055909156799], [1908.396118581295, 1450.1928869485855, 0.9719980955123901], [1901.281094610691, 1473.2697660923004, 0.9568612575531006], [1902.3744297027588, 1490.7288138866425, 0.9535452127456665], [1919.6365355849266, 1471.8723105192184, 0.8522146940231323], [1917.1035763025284, 1497.4250456094742, 0.9186901450157166], [1915.799398124218, 1505.544674873352, 0.9714179039001465], [1915.9174289107323, 1506.9655919075012, 0.9617753028869629], [1938.335991859436, 1470.4997633695602, 0.8432860374450684], [1934.964259505272, 1496.228717803955, 0.9244122505187988], [1933.021092414856, 1501.7569408416748, 0.9856633543968201], [1932.1070773005486, 1500.709546804428, 0.9512953758239746], [1954.1973865032196, 1467.8032456636429, 0.906885027885437], [1952.3831540942192, 1487.333206653595, 0.9238507747650146], [1950.0320469737053, 1492.6881042718887, 0.9762418270111084], [1947.7674992084503, 1492.8244287967682, 0.9661561250686646], [1964.6534465551376, 1466.6316195726395, 0.957784116268158], [1964.3601266145706, 1479.4636951684952, 0.9727674722671509], [1962.552922308445, 1484.094698190689, 1.0022079944610596], [1960.1001144051552, 1485.4953577518463, 0.9795989990234375], [1704.9451496005058, 1329.7077422142029, 0.8573814630508423], [1706.8013955950737, 1350.515474319458, 0.8936662673950195], [1707.7446083426476, 1366.1250975131989, 0.8643858432769775], [1706.213205397129, 1387.6610651016235, 0.8806315064430237], [1706.2636423707008, 1405.8051462173462, 0.8623154759407043], [1685.5378703474998, 1388.2433640956879, 0.8812478184700012], [1689.8018617033958, 1409.5157778263092, 0.9175753593444824], [1698.6052323579788, 1420.1517779827118, 0.8866860270500183], [1704.4083782434464, 1425.6540590524673, 0.8614552617073059], [1687.2092669010162, 1392.4654968976974, 0.9244275689125061], [1693.7211142778397, 1412.340868473053, 0.9598500728607178], [1703.9385805726051, 1420.1069221496582, 0.903150737285614], [1710.135300219059, 1422.6473746299744, 0.8694174885749817], [1693.3942041397095, 1391.2624510526657, 0.9386492371559143], [1698.6404245495796, 1408.2294282913208, 0.9550232291221619], [1707.3110880851746, 1414.6469130516052, 0.9241269826889038], [1712.7438632845879, 1417.0038080215454, 0.8950108885765076], [1698.7786611318588, 1389.511213183403, 0.9072856903076172], [1702.5424169301987, 1403.0161262750626, 0.917733371257782], [1707.6906056404114, 1409.0185395479202, 0.9306231737136841], [1711.7040657997131, 1411.7252893447876, 0.9123875498771667]]
    print(len(kps_in))
    print(cocowb_2_openpose(kps_in))